import argparse
import math
import time
import warnings
from contextlib import nullcontext

import pandas as pd
import torch
import torch.distributed as dist
from LMConfig import LMConfig
from dataset import PretrainDataset
from model import Transformer
from torch import optim
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
import os

warnings.filterwarnings('ignore')  # 忽略警告信息


# 定义日志打印函数
def Logger(content):
    if not ddp or dist.get_rank() == 0:  # 如果不是分布式训练或当前进程是主进程
        print(content)


# 学习率调度函数
def get_lr(it, all):
    warmup_iters = args.warmup_iters  # 预热迭代次数
    lr_decay_iters = all  # 总迭代次数
    min_lr = args.learning_rate / 10  # 最小学习率

    if it < warmup_iters:  # 如果在预热阶段
        return args.learning_rate * it / warmup_iters  # 线性增加学习率
    if it > lr_decay_iters:  # 如果超过总迭代次数
        return min_lr  # 返回最小学习率
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)  # 计算衰减比例
    assert 0 <= decay_ratio <= 1  # 确保衰减比例在合理范围内
    coeff = 0.5 * (1. + math.cos(math.pi * decay_ratio))  # 余弦衰减公式
    return min_lr + coeff * (args.learning_rate - min_lr)  # 返回当前学习率


# 训练一个epoch的函数
def train_epoch(epoch, wandb: None):
    start_time = time.time()  # 记录开始时间
    for step, (X, Y, loss_mask) in enumerate(train_loader):  # 遍历训练数据
        X, Y, loss_mask = X.to(args.device), Y.to(args.device), loss_mask.to(args.device)  # 将数据移动到指定设备

        lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        with ctx:  # 混合精度上下文
            out = model(X, Y)  # 前向传播
            loss = out.last_loss / args.accumulation_steps  # 计算损失并归一化
            loss_mask = loss_mask.view(-1)  # 展平损失掩码
            loss = torch.sum(loss * loss_mask) / loss_mask.sum()  # 计算加权损失

        scaler.scale(loss).backward()

        if (step + 1) % args.accumulation_steps == 0:  # 如果达到梯度累积步数
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

            scaler.step(optimizer)
            scaler.update()

            optimizer.zero_grad(set_to_none=True)

        if step % args.log_interval == 0:  # 每隔一定步数打印日志
            spend_time = time.time() - start_time  # 计算耗时
            Logger('Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
                epoch, args.epochs, step, iter_per_epoch, loss.item() * args.accumulation_steps,
                optimizer.param_groups[-1]['lr'], spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))

        if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):  # 每隔一定步数保存模型
            model.eval()  # 切换到评估模式
            moe_path = '_moe' if lm_config.use_moe else ''  # 判断是否使用MoE
            ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}.pth'  # 保存路径

            if isinstance(model, torch.nn.parallel.DistributedDataParallel):  # 如果是分布式模型
                state_dict = model.module.state_dict()  # 获取模型参数
            else:
                state_dict = model.state_dict()

            torch.save(state_dict, ckp)  # 保存模型
            model.train()  # 切换回训练模式
    return model


# 初始化模型函数
def init_model():
    def count_parameters(model):  # 计算模型参数量
        return sum(p.numel() for p in model.parameters() if p.requires_grad)

    tokenizer = AutoTokenizer.from_pretrained('model')  # 加载预训练分词器
    model = Transformer(lm_config).to(args.device)  # 初始化模型并移动到指定设备

    Logger(f'LLM总参数量：{count_parameters(model) / 1e6:.3f} M')  # 打印模型参数量
    return model, tokenizer


# 主函数
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="miniDeepSeek Pretraining")  # 定义命令行参数解析器
    parser.add_argument("--out_dir", type=str, default="out", help="Output directory")  # 输出目录
    parser.add_argument("--epochs", type=int, default=2, help="Number of epochs")  # 训练轮数
    parser.add_argument("--batch_size", type=int, default=1, help="Batch size")  # 批量大小
    parser.add_argument("--learning_rate", type=float, default=2e-4, help="Learning rate")  # 学习率
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="Device to use")  # 设备
    parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")  # 数据类型
    parser.add_argument("--use_wandb", action="store_true", help="Use Weights & Biases")  # 是否使用WandB
    parser.add_argument("--wandb_project", type=str, default="miniDeepSeek-Pretrain", help="Weights & Biases project name")  # WandB项目名称
    parser.add_argument("--num_workers", type=int, default=1, help="Number of workers for custom_data loading")  # 数据加载线程数
    parser.add_argument("--data_path", type=str, default="dataset/pretrain_data.csv", help="Path to training custom_data")  # 数据路径
    parser.add_argument("--ddp", action="store_true", help="Use DistributedDataParallel")  # 是否使用分布式训练
    parser.add_argument("--accumulation_steps", type=int, default=8, help="Gradient accumulation steps")  # 梯度累积步数
    parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")  # 梯度裁剪阈值
    parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")  # 预热迭代次数
    parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")  # 日志打印间隔
    parser.add_argument("--save_interval", type=int, default=1000, help="Model saving interval")  # 模型保存间隔
    parser.add_argument('--local_rank', type=int, default=-1, help='local rank for distributed training')  # 分布式训练的本地进程编号
    args = parser.parse_args()  # 解析命令行参数

    lm_config = LMConfig()  # 加载语言模型配置
    max_seq_len = lm_config.max_seq_len  # 最大序列长度

    args.save_dir = os.path.join(args.out_dir)  # 设置保存目录
    os.makedirs(args.save_dir, exist_ok=True)  # 创建保存目录
    os.makedirs(args.out_dir, exist_ok=True)  # 创建输出目录

    tokens_per_iter = args.batch_size * max_seq_len  # 每次迭代的token数量
    torch.manual_seed(1337)  # 设置随机种子
    device_type = 'cuda' if torch.cuda.is_available() else 'cpu'  # 判断设备类型
    ctx = nullcontext() if device_type == 'cpu' else torch.cuda.amp.autocast()  # 设置混合精度上下文

    ddp = False  # 是否使用分布式训练

    model, tokenizer = init_model()  # 初始化模型和分词器

    df = pd.read_csv(args.data_path)  # 加载训练数据
    df = df.sample(frac=0.1)  # 随机采样10%的数据
    train_ds = PretrainDataset(df, tokenizer, max_length=max_seq_len)  # 创建训练数据集
    train_sampler = None  # 数据采样器
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch_size,
        pin_memory=False,
        drop_last=False,
        shuffle=False,
        num_workers=args.num_workers,
        sampler=train_sampler,
    )  # 创建数据加载器

    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['bfloat16', 'float16']))  # 初始化混合精度缩放器
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)  # 初始化优化器

    iter_per_epoch = len(train_loader)  # 每个epoch的迭代次数
    for epoch in range(args.epochs):  # 遍历每个epoch
        train_epoch(epoch, None)  # 训练一个epoch
        model.eval()  # 切换到评估模式
        moe_path = '_moe' if lm_config.use_moe else ''  # 判断是否使用MoE
        ckp = f'{args.save_dir}/pretrain_{lm_config.dim}{moe_path}2.pth'  # 保存路径
        state_dict = model.state_dict()
        torch.save(state_dict, ckp)  # 保存模型
        model.train()  # 切换回训练模式